#!/usr/bin/env python
'''
main function of the program
'''

from __future__ import division
import numpy
import random
import ConfigParser
import sys
import os
import math
import time
import lattice
import cPickle as pickle

def argument(input):
    '''
    if input in sys.argv returns argument of input
    else returns False
    '''
    if input in sys.argv:
        try:
            return sys.argv[ sys.argv.index(input) + 1 ]
        except:
            return True
    else:
        return False
        
def choice(path):
    '''
    path - path to the directory in which we are interested in. Function analyzes what subdirectories are in our directory
    and if there is more then one, asks user in which of them we are interested. As a result we have path to subdirectory    
    '''
    if type(path) == list:
        fList = path
    else:
        fList = os.listdir( path )
    if len(fList) == 0:
        raise 'Zero variants. No choice'
    elif len(fList) == 1:
        return fList[0]
    else:
        for i, case in enumerate( fList ):
            print i,  ' - ',  case
        caseNumber = int(raw_input('Enter number of the case we are interested in\n'))
        return fList[ caseNumber]


if '-cName' in sys.argv:
    computername = sys.argv[ sys.argv.index('-cName') + 1 ] #name of the computer we are starting at
else:
    computername = 'local'

B = eval(argument('-B'))

if type(B) == float or type(B) == int:
    bRange = [ B ]
else:
    bRange = B
    
config = ConfigParser.RawConfigParser()
config.read('config.py') #opening configuration file

if '-m' in sys.argv:
    modelname = 'Triangular_transverse_2d'
elif __name__ != '__main__':
    modelname = 'Triangular_transverse_2d'
else:
    ModelNames = config.sections()
    
    modelname = choice( ModelNames )

if __name__ == '__main__':
    starttime = time.time()

h_shape = eval(config.get (modelname, 'horizontal shape' ))
if '-N' in sys.argv:
    N = int( argument('-N') )
    h_shape = tuple( [N] * len(h_shape) )


max_T = config.getfloat( modelname, 'maximum temperature')
min_T = config.getfloat( modelname, 'minimum temperature')
step_T = config.getfloat( modelname, 'temperature step')

if min_T != max_T:
    tRange = numpy.arange( min_T, max_T + step_T, step_T )
else:
    tRange = [ min_T ]

h_neighbours = eval(config.get( modelname, 'horizontal neighbours'))

ferromagnetic = config.getboolean(modelname, 'ferromagnetic')

try:
    v_neighbours = eval(config.get( modelname, 'vertical neighbours'))
except:
    v_neighbours = ()


Sweeps = config.getint(modelname, 'sweeps')
Equilib = config.getint(modelname, 'equlib' )

#working with time, for proper to file saving
fPrefix = time.strftime('%Y_%m_%d_%H_%M', time.gmtime())

if ferromagnetic:
    J = 1
else:
    J = -1
    
dt = 1 / 8

def main(T,  B):
    '''
    main function of the program
    '''
    
    #we have two different cases. We can consider classical model in which we vary T or 
    #we can consider quantum case and vary B
    
    #we need to define shape. It looks a bit tricky. Now I assume that vertical size is included in shape
    if not argument('-q'):#Considering classic case
        if B == 0:
            v_size = 1
        else:
            v_size = max( int( 1 / ( dt * T ) ) + 1, 2)
    
        dT = 1 / ( v_size * T)
        
    else:#if we are considering quantum case
        h_size = h_shape[0]
        v_size = 2 * h_size #this 2 is not some random number - it is some magical stuff from Richard's experience
        dT = dt
        
        T = 1 / (v_size * dT)        
    
    shape = tuple( [v_size] + list(h_shape) )
    
    if B == 0 or not v_neighbours:
        LD = 0
    else: 
        LD = - 0.5 * math.log( math.tanh( B * dT ) )
    
    if '-info' in sys.argv:
        print
        print 'shape - ', shape
        print 'temperature - ', T
        print 'B - ', B
        print 'Delta Tau - ', dT
        print 'Lambda - ',  LD

    #calculate number of sites
    nSites = numpy.prod(shape)
    
    #Generate Lattice    
    tLattice = lattice.Lattice(h_shape, v_size, h_neighbours, ferromagnetic, v_neighbours = v_neighbours, LD = LD, T = T, B = B)
    
    #Calculate initial Energies
    if B != 0:
        E_J = tLattice.Energy_J()
        
    E_B = tLattice.Energy_B()
    ES_J = []#list with E_J energies
    ES_B = []#list with E_B energies
    MS = []
    acceptence = 0
    norm_acc = 0 # normalization for acceptence
    
    def GrowCluster(index,  ClusterSpin ):
        tLattice.lattice[ index ] *= -1                
        
        for neighbor in tLattice.h_neighbours[ index ]:#we are going through all horizontal neighbors of our site            
            if tLattice.lattice[ neighbor ] == ClusterSpin:#we are working only with sites that are aligned to our spin
                if cluster[ neighbor ] == 0:#if this site not in cluster
                    weight[ neighbor ] = 1 #we count this try
                    tryAdd_horizontal( neighbor, ClusterSpin )#We try add this site to cluster
                    
        for neighbor in tLattice.v_neighbours[ index ]:#we are going through all vertical neighbors of our site            
            if tLattice.lattice[ neighbor ] == ClusterSpin:#we are working only with sites that are aligned to our spin
                if cluster[ neighbor ] == 0:#if this site not in cluster
                    weight[ neighbor ] = 1 #we count this try
                    tryAdd_vertical( neighbor, ClusterSpin )#We try add this site to cluster            
        
                
    def tryAdd_horizontal( index, ClusterSpin):
        '''
        function tries to add site to cluster
        '''                        
        if random.random() < ( 1 - math.exp(  - 2 * J * dT) ): 
            GrowCluster( index,  ClusterSpin )
            
    def tryAdd_vertical( index, ClusterSpin):
        '''
        function tries to add site to cluster
        '''                        
        if random.random() < ( 1 - math.exp(  - 2 * LD ) ): 
            GrowCluster( index,  ClusterSpin )        
        
    
    for sweep in xrange(Sweeps):         
                       
        cluster = numpy.zeros( shape )#sites included to cluster                        
        weight = numpy.zeros( shape )#counts spins we tried to flip
        
        index = tLattice.RandomSite()
        cluster[ index ] = 1#we include our site to cluster
        weight[ index ] = 1#we include it at the first try
        
        ClusterSpin = tLattice.lattice[ index ]#Spin of our cluster
        
        GrowCluster(index,  ClusterSpin )#building and flipping cluster
        
        t_clust = numpy.sum( cluster )#number of the flipped spins
        t_norm = numpy.sum( weight )#number of the tries to flip
        
        if t_clust > t_norm:
            raise 'number of sites in cluster should be less then number of considered sites'
                        
        acceptence += numpy.sum( cluster )
        
        norm_acc += numpy.sum( weight )
                    
        if sweep > Equilib and sweep % 10 == 0: #We save only each tenth sweep result
            if '-q' in sys.argv:#We are considering quantum case
                tM = numpy.sum( tLattice.lattice )
            else:#We are considering classic case
                tM = tLattice.sum_layer(0)#We are storing magnetization of the one layer
            MS += [ tM ]
            ES_J += [ tLattice.Energy_J() ]
            
            if B == 0:                
                ES_B += [ 0 ]
            else:                
                ES_B += [ tLattice.Energy_B() ]
                
    if Sweeps > Equilib:
        #small check for energy consistency recomment when something is suspecios
        """
        if abs(E_B - tLattice.Energy_B()) >  0.00000001:
            raise 'Something is wrong with E_B'
        if abs(E_J - tLattice.Energy_J()) >  0.00000001:
            raise 'Something is wrong with E_J'
        """
        if B == 0:
            acceptence /= norm_acc
        else:
            acceptence /= (Sweeps * nSites)
        if '-q' in sys.argv:
            fName = open('results/%s/%s/Q_%s_%s_B%s_T%s'%(modelname, case, fPrefix, str(shape), str(B), str(T)),'w')
        else:
            fName = open('results/%s/%s/%s_%s_B%s_T%s'%(modelname, case, fPrefix, str(shape), str(B), str(T)),'w')
        pickle.dump({'shape': shape[1:],
                     'L': shape[0],
                     'T': T,
                     'E_J': ES_J,
                     'E_B': ES_B,
                     'M': MS,
                     'Sweeps': Sweeps,
                     'Equilib': Equilib,
                     'B': B,
                     'acceptence': acceptence
                     }, 
                     fName)        
        fName.close()
        if '-info' in sys.argv:
            print 'acceptence', acceptence
        return 
    
        
    else:
        return 0, 0, 0, 0

if __name__ == "__main__":

    try:
        os.mkdir('results')
    except:
        pass

    try:
        os.mkdir(os.path.join( 'results', modelname ) )
    except:
        pass
    
    if ferromagnetic:
        case = 'ferromagnetic'
    else:
        case = 'antiferromagnetic'
    
    try:
        os.mkdir( os.path.join('results', modelname, case ) )
    except:
        pass

    for T in tRange:
        for B in bRange:
            main(T,  B)
    working_time = time.time() - starttime
    
    fName = open('{0}_finished'.format(computername), 'w')
    print >> fName, working_time
    fName.close()
    if '-info' in sys.argv:
        print 'program worked for', working_time, ' seconds'
